//////////////////////////////////////////////////////////////////
// MattiseOS Kernel
// Date: 2008
// Author: Matthew Iselin
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
// 
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//////////////////////////////////////////////////////////////////

// Includes
#include <mattise.h>
#include <stdint.h>
#include <mem/page.h>
#include <mem/util.h>
#include <mem/heap.h>
#include <kernio/io.h>

// pretty self-explanatory
#define USEBITMAP
//#define USESTACK

// allocates contiguous space and maps it in to the parent page directory
uint32_t kAllocateAndMap( uint32_t count, uint32_t vaddr, uint32_t flags )
{
	// nice and simple
	return kAllocateAndMapPageDir(0xfffff000, count, vaddr, flags);
}

// allocates contiguous space and maps it in to a specific page dir
uint32_t kAllocateAndMapPageDir( uint32_t pagedir, uint32_t count, uint32_t vaddr, uint32_t flags )
{
	uint32_t i;
	vaddr &= ~0xFFF;
	for(i = 0; i < count; i++)
	{
		// get a page
		uint32_t* tmp = kAllocatePages(1);
		//dprintf("mapping for %x, tmp = %x\n", vaddr + (i * 0x1000), tmp);
		//uint32_t addr = (uint32_t) tmp;

		// map it in - this makes non-contiguous pages appear contiguous, which is nice
		MapPage(pagedir, vaddr + (i* 0x1000), (uint32_t) tmp, flags);
	}
	return vaddr;
}

// allocates contiguous space and maps it into two seperate address spaces
uint32_t kAllocateAndMapPageDirTwice( uint32_t pagedir1, uint32_t pagedir2, uint32_t count, uint32_t vaddr1, uint32_t vaddr2, uint32_t flags1, uint32_t flags2 )
{
	uint32_t i;
	vaddr1 &= ~0xFFF;
	vaddr2 &= ~0xFFF;
	// 7540736 -> 7730541
	for(i = 0; i < count; i++)
	{
		// get a page
		uint32_t* tmp = kAllocatePages(1);
		//uint32_t addr = (uint32_t) tmp;
#if DEBUG
		dprintf("mapping %x,%x to %x\n", vaddr1 + (i * 0x1000), vaddr2 + (i * 0x1000), tmp);
#endif

		// map it in - this makes non-contiguous pages appear contiguous, which is nice
		MapPage(pagedir1, vaddr1 + (i* 0x1000), (uint32_t) tmp, flags1);
		MapPage(pagedir2, vaddr2 + (i* 0x1000), (uint32_t) tmp, flags2);
	}
	return vaddr1;
}

// frees pages mapped with the previous functions
void kFreeMappedPages( uint32_t count, uint32_t vaddr )
{
	uint32_t i;
	vaddr &= ~0xFFF;
	for(i = 0; i < count; i++)
	{
		// get the physical mapping
		uint32_t phys = GetPhysicalAddress( 0xFFFFF000, vaddr );
		kFreePages( (uint32_t*) phys );
	}
}

#ifdef USESTACK

// Variables

// the global memory stack
uint32_t memstack[0xfffff] = {0}; // can keep track of a 32-bit address space

// pointer to the stack - using a signed int because it's easier to detect
// problems, and the stack index fits into a signed int
int32_t stack_ptr = 0;

// the number of free pages in the system
uint32_t num_free_pages = 0;

// Functions

// sets a set of pages to either on or off
void kSetPages( uint32_t start, uint32_t end, uint32_t set )
{
	// this is a bit of a hacky way to do it but it works
	uint32_t i;
	start &= ~0xFFF;
	end &= ~0xFFF;
	if(set == 0)
		for(i = end; i >= start; i -= 0x1000)
			kFreePages((uint32_t*) i);
}

// allocates one or more contiguous pages
uint32_t* kAllocatePages( uint32_t UNUSED(count) )
{
	// pop the address
	if((stack_ptr - 1) < 0)
	{
		//dprintf("out of memory - %d!\n", stack_ptr);
		return 0; // no free memory
	}
	uint32_t ret = memstack[stack_ptr--];
	if(ret == 0)
		dprintf("==== [WARN ] kAllocatePages returning 0 - stack_ptr = %d/%x! ====\n", stack_ptr, stack_ptr);
	num_free_pages--;
	return (uint32_t*) ret;
}

// frees pages starting at a specific address
void kFreePages( uint32_t* addr )
{
	// push the address onto the stack
	if(addr == 0)
		dprintf("==== [WARN ] kFreePages with null page ====\n");
	memstack[++stack_ptr] = ((uint32_t) addr) & ~0xFFF;
	if(stack_ptr >= 0xfffff)
		stack_ptr %= 0xfffff;
	num_free_pages++;
}

// finds out how much memory is available
uint32_t kGetAvailMem()
{
	return num_free_pages * 0x1000;
}

// sets up the page allocator
void kInitPageAlloc( uint32_t UNUSED(maxmem) )
{
    // set the pointer to the right location?
	stack_ptr = 0;
}


#elif defined(USEBITMAP)

// Defined in kheap.c
//extern uint32_t placement_address;

// Macros used in the bitset algorithms.
#define INDEX_FROM_BIT(a) (a/(8*4))
#define OFFSET_FROM_BIT(a) (a%(8*4))

// A bitset of frames - used or free.
uint32_t frames[INDEX_FROM_BIT(0xfffff000 / 0x1000)] = {0};
uint32_t nframes;

// the number of free pages
uint32_t num_free_pages = 0;

// Static function to set a bit in the frames bitset
static void set_frame(uint32_t frame_addr)
{
   uint32_t frame = frame_addr/0x1000;
   uint32_t idx = INDEX_FROM_BIT(frame);
   uint32_t off = OFFSET_FROM_BIT(frame);
   frames[idx] |= (0x1 << off);
}

// Static function to clear a bit in the frames bitset
static void clear_frame(uint32_t frame_addr)
{
   uint32_t frame = frame_addr/0x1000;
   uint32_t idx = INDEX_FROM_BIT(frame);
   uint32_t off = OFFSET_FROM_BIT(frame);
   frames[idx] &= ~(0x1 << off);
}

// Static function to test if a bit is set.
/*static uint32_t test_frame(uint32_t frame_addr)
{
   uint32_t frame = frame_addr/0x1000;
   uint32_t idx = INDEX_FROM_BIT(frame);
   uint32_t off = OFFSET_FROM_BIT(frame);
   return (frames[idx] & (0x1 << off));
}*/

// Static function to find the first free frame.
static uint32_t first_frame()
{
   uint32_t i, j;
   for (i = 0; i < INDEX_FROM_BIT(nframes); i++)
   {
       if (frames[i] != 0xFFFFFFFF) // nothing free, exit early.
       {
           // at least one bit is free here.
           for (j = 0; j < 32; j++)
           {
               uint32_t toTest = 0x1 << j;
               if ( !(frames[i]&toTest) )
               {
                   return i*4*8+j;
               }
           }
       }
   }
   return 0;
}

// sets a set of pages to either on or off
void kSetPages( uint32_t start, uint32_t end, uint32_t set )
{
	// this is a bit of a hacky way to do it but it works
	uint32_t i;
	start &= ~0xFFF;
	end &= ~0xFFF;
	if(set == 1)
	{
		for(i = start; i <= end; i += 0x1000)
		{
			set_frame(i);
			num_free_pages--; // -= 0x1000;
		}
	}
}

// allocates one or more contiguous pages
uint32_t* kAllocatePages( uint32_t UNUSED(count) )
{
	uint32_t ret = first_frame() * 0x1000;
	if(ret == 0)
		return 0;
	set_frame(ret);
	num_free_pages--; // -= 0x1000;
	return (uint32_t*) ret;
}

// frees pages starting at a specific address
void kFreePages( uint32_t* addr )
{
	if(addr == 0)
		return;
	num_free_pages++; // += 0x1000;
	clear_frame(((uint32_t) addr));
}

// finds out how much memory is available
uint32_t kGetAvailMem()
{
	return num_free_pages * 0x1000;
}

// sets up the page allocator
void kInitPageAlloc( uint32_t maxmem )
{
	nframes = maxmem / 0x1000;
	num_free_pages = nframes;
}

#if 0
// Defines

// number of bits in a bitmap entry
#define NUMBITS			32

// total number of bytes in the bitmap
#define NUMBITMAPBYTES	((0xffffffff/0x1000/NUMBITS)+1)

// the type name
#define TYPENAME		uint32_t

// bitset stuff
#define BIT( b )			( 1 << b )
#define SETBIT( b, d )		( d |= BIT( b ) )
#define RESETBIT( b, d )	( d &= (~BIT( b )) )

// Variables

// the global memory bitmap
TYPENAME membitmap[NUMBITMAPBYTES];

// the global memory index (see the design documentation)
TYPENAME memindex[NUMBITMAPBYTES];

// the location of the last free page (to save time)
uint32_t lastfreepage = 0;

// last valid page
uint32_t lastvalid = 0;

// a count of how many pages are free (makes getting available memory fast)
uint32_t freepagecount = 0;

// Functions

// sets a set of pages to either on or off
void kSetPages( uint32_t start, uint32_t end, uint32_t set )
{
    // set them all
    uint32_t i = start / 0x1000;
    for( ; i < (end / 0x1000); i++ )
    {
        // figure out what to do
        if( set )
            membitmap[i/NUMBITS] |= (1 << (i % NUMBITS));
        else
            membitmap[i/NUMBITS] &= ~(1 << (i % NUMBITS));
    }
    if( set )
		freepagecount -= (end - start) / 0x1000;
}

// frees pages starting at a specific address
void kFreePages( uint32_t* addr )
{
	// valid address?
	if( addr == 0 )
		return; // zero is *NEVER* meant to be allocated

    // convert the pointer argument into an integer
    uint32_t intaddr = (uint32_t) addr;
    uint32_t i = intaddr / 0x1000;

	// the number of freed pages
	int nFreedPages = 0;

    // we work off the index table here
    while( 1 )
    {
		// check that this bit is set in the index table
		if( memindex[i/NUMBITS] & ( 1 << ( i % NUMBITS ) ) )
		{
			// yes, so clear it and then clear the bitmap bit
			memindex[i/NUMBITS] &= ~( 1 << ( i % NUMBITS ) );
			membitmap[i/NUMBITS] &= ~( 1 << ( i % NUMBITS ) );

			// another free page!
			freepagecount++;
			nFreedPages++;
        }
        else
        {
            // we hit a 0 - we *MUST* still free it though (0 = end of chain)
            membitmap[i/NUMBITS] &= ~( 1 << ( i % NUMBITS ) );

			// another free page!
			freepagecount++;
			nFreedPages++;

            // set the last free page to use memory efficiently
            lastfreepage = 0; //intaddr / 0x1000;

            // return to the caller
            return;
        }

        // go to the next page
        i++;
    }
}

// allocates one or more contiguous pages
uint32_t* kAllocatePages( uint32_t count )
{
    // the start bit of the set of pages - null is not allowed!
    uint32_t set_startbit = 0;
    uint32_t set_count = 0;

    // try to find a free bit in the bitmap
    uint32_t i;
    for( i = lastfreepage; i < (NUMBITMAPBYTES*NUMBITS); i++ )
    {
    	// check if we should skip
		if( membitmap[i / NUMBITS] == 0xffffffff && ( i % NUMBITS ) == 0 )
    	{
			i += NUMBITS;
			continue;
    	}

        // is this bit set?
        if( membitmap[i / NUMBITS] & ( 1 << ( i % NUMBITS ) ) )
        {
            // yes, so stop any group so far
            set_startbit = 0;
            set_count = 0;
        }
        else
        {
            // if there is a page to start from, increment the counter
            if( set_startbit > 0 )
                set_count++;
            else // (otherwise, set the starter)
            {
                set_startbit = i;
                set_count = 1;

                // the last known free page is here
                lastfreepage = 0; //i;
            }

            // check if we've hit the count
            if( set_count == count )
            {

				// these are now being used
				freepagecount -= set_count;

                // set each bit as taken
                for( i = 0; i < set_count; i++ )
                {
                    // set the index list element (unless we're almost done, in which case, clear it)
                    if( i + 1 == set_count )
                        memindex[(set_startbit+i)/NUMBITS] &= ~( 1 << ( ( set_startbit + i ) % NUMBITS ) );
                    else
                        memindex[(set_startbit+i)/NUMBITS] |= ( 1 << ( ( set_startbit + i ) % NUMBITS ) );

                    // and then set the actual memory up
                    membitmap[(set_startbit+i)/NUMBITS] |= ( 1 << ( ( set_startbit + i ) % NUMBITS ) );
                }
                return (uint32_t*) (set_startbit * 0x1000); // return it
            }
        }
    }

    // none found
    return (uint32_t*) 0;
}

// finds out how much memory is available
uint32_t kGetAvailMem()
{
	return freepagecount << 12;
}

// sets up the page allocator
void kInitPageAlloc( uint32_t maxmem )
{
    // zero out the bitmap
    memset( membitmap, 0, NUMBITMAPBYTES * sizeof( TYPENAME ) );

    // clear out the memory index
    memset( memindex, 0, NUMBITMAPBYTES * sizeof( TYPENAME ) );

    // set the maxmimum accessible address
    lastvalid = freepagecount = maxmem;
    freepagecount /= 0x1000;
}

#endif

#else
#error "NO PAGE MANAGER!"
#endif
